{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Continued Pre-Training (CPT) at scale with DeepSpeed\n",
        "\n",
        "This guide provides a step-by-step workflow for continued pre-training the [`google/gemma-3-4b-pt`](https://huggingface.co/google/gemma-3-4b-pt) model on a multi-GPU Anyscale cluster. It uses LLaMA-Factory for the training framework and `DeepSpeed` to efficiently manage memory and scale the training process.\n",
        "\n",
        "CPT is a technique to further adapt a pre-trained base model on large-scale unlabeled text. By continuing to train on high-quality corpora, you adapt the model to new domain knowledge and improve generalization. This notebook performs full fine-tuning of the base model instead of using parameter-efficient fine-tuning (PEFT) techniques.\n",
        "\n",
        "- **Full fine-tuning vs LoRA:** Full fine-tuning generally yields the best quality but requires significantly more compute, longer training, and large checkpoints. LoRA is much faster and cheaper with small adapter checkpoints, but typically shows the most improvement on curated, simplified corpora (gains on broad/noisy corpora may be limited). See [Compare full vs freeze vs PEFT](https://docs.anyscale.com/llm/fine-tuning#compare-full-vs-freeze-vs-parameter-efficient-fine-tuning-peft) and [LoRA speed and memory optimizations](https://docs.anyscale.com/llm/fine-tuning/speed-and-memory-optimizations#lora).\n",
        "\n",
        "## Step 1: Set up your environment\n",
        "\n",
        "### Dependencies\n",
        "First, ensure your environment has the correct libraries. Start with a pre-built container image and install LLaMA-Factory and DeepSpeed on top of it.\n",
        "\n",
        "Recommended container image:\n",
        "```bash\n",
        "anyscale/ray-llm:2.48.0-py311-cu128\n",
        "```\n",
        "\n",
        "Execute the following commands to install the required packages and optional tools for experiment tracking and faster model downloads:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%%bash\n",
        "# Install the specific version of LLaMA-Factory\n",
        "pip install -q llamafactory==0.9.3\n",
        "\n",
        "# Install DeepSpeed for large-scale training\n",
        "pip install -q deepspeed==0.16.9\n",
        "\n",
        "# (Optional) For accelerated model downloads from Hugging Face\n",
        "pip install -q hf_transfer==0.1.9\n",
        "\n",
        "# (Optional) Experiment tracking library\n",
        "pip install -q mlflow==3.4.0\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Model and compute resources\n",
        "\n",
        "DeepSpeed ZeRO-3 partitions parameters, gradients, and optimizer states across multiple GPUs, enabling CPT of mid-sized LLMs on just 4 GPUs.\n",
        "\n",
        "| Item | Value |\n",
        "|------|-------|\n",
        "| **Base model** | [`google/gemma-3-4b-pt`](https://huggingface.co/google/gemma-3-4b-pt) |\n",
        "| **Worker nodes** | 4 × L40S / 4 x A100-40G |\n",
        "\n",
        "## Step 2: Prepare the dataset\n",
        "\n",
        "### Understand the dataset\n",
        "This tutorial uses a simple JSONL corpus ([C4](https://huggingface.co/datasets/allenai/c4)) containing cleaned English web text derived from Common Crawl, widely used for language-model pretraining. Each line is a JSON object with at least a `text` field. For demo purposes, the sample `c4.jsonl` contains only the first 100 records from the original C4 dataset (hosted on S3) to enable quick runs.\n",
        "\n",
        "**Dataset example**\n",
        "\n",
        "```json\n",
        "{\"text\": \"Beginners BBQ Class Taking Place in Missoula!\\nDo you want to get better at making delicious BBQ? You will have the opportunity, put this on your calendar now. Thursday, September 22nd join World Class BBQ Champion, Tony Balay from Lonestar Smoke Rangers. He will be teaching a beginner level class for everyone who wants to get better with their culinary skills.\\nHe will teach you everything you need to know to compete in a KCBS BBQ competition, including techniques, recipes, timelines, meat selection and trimming, plus smoker and fire information.\\nThe cost to be in the class is $35 per person, and for spectators it is free. Included in the cost will be either a t-shirt or apron and you will be tasting samples of each meat that is prepared.\", \"timestamp\": \"2019-04-25 12:57:54\", \"url\": \"https://klyq.com/beginners-bbq-class-taking-place-in-missoula/\"}\n",
        "```\n",
        "\n",
        "### Register the dataset\n",
        "\n",
        "To specify new datasets that are accessible across Ray worker nodes, you must first add a **`dataset_info.json`** to **[storage shared across nodes](https://docs.anyscale.com/configuration/storage#shared)** such as `/mnt/cluster_storage`. This configuration file acts as a central registry for all your datasets. It maps a custom name to your dataset file location, format, and column structure. \n",
        "\n",
        "If you plan to run CPT on this text dataset, first complete the setup steps below. Ensure that you place the dataset files in a storage location that all workers can access (for example, a shared mount or object storage). Avoid storing large files on the head node.\n",
        "\n",
        "`dataset_info.json`\n",
        "```json\n",
        "{\n",
        "  \"my_cpt_c4\": {\n",
        "      \"file_name\": \"/mnt/cluster_storage/c4.jsonl\",\n",
        "      \"columns\": {\n",
        "          \"prompt\": \"text\"\n",
        "      }\n",
        "  }\n",
        "}\n",
        "```\n",
        "\n",
        "For a more detailed dataset preparation and formatting guide, see [Choose your data format](https://docs.anyscale.com/llm/fine-tuning/data-preparation#continued-pretraining).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%%bash\n",
        "# Make sure all files are accessible to worker nodes\n",
        "# Create a copy of the data in /mnt/cluster_storage\n",
        "wget https://anyscale-public-materials.s3.us-west-2.amazonaws.com/llm-finetuning/llama-factory/datasets/alpaca/c4.jsonl -O /mnt/cluster_storage/c4.jsonl\n",
        "# Create a copy of the dataset registry in /mnt/cluster_storage\n",
        "cp ../dataset-configs/dataset_info.json /mnt/cluster_storage/\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 3: Create the pre-training config (CPT with DeepSpeed)\n",
        "\n",
        "Next, create the main YAML configuration file—the master recipe for your pre-training job. It specifies the base model, the training method (full fine-tuning), the dataset, training hyperparameters, cluster resources, and more.\n",
        "\n",
        "**Important notes:**\n",
        "- **MLflow tracking:** To track experiments with MLflow, set `report_to: mlflow` in the config. If you don't want to use MLflow, set `report_to: none` to avoid errors.\n",
        "- **Access and paths:** The YAML only needs to be on the **head node**, but any referenced paths (`dataset_dir`, `output_dir`) must reside on storage **reachable by all workers** (for example, `/mnt/cluster_storage/`).\n",
        "- **Gated models:** If your base model has gated access (for example, Gemma) on Hugging Face, set `HF_TOKEN` in the runtime environment.\n",
        "- **GPU selection and placement:** The config uses a 4xL40S node (`anyscale/accelerator_shape:4xL40S`) so that all 4 GPUs are on the same machine, which is important for efficient DeepSpeed ZeRO-3 communication. You can switch to other multi-GPU nodes such as `4xA100-40GB` or any other node type with comparable or more VRAM, depending on your cloud availability.\n",
        "\n",
        "### Configure LLaMA-Factory with Ray\n",
        "\n",
        "**Note**: To customize the training configuration, edit `train-configs/cpt_deepspeed.yaml`. \n",
        "\n",
        "```yaml\n",
        "# cpt_deepspeed.yaml\n",
        "\n",
        "### model\n",
        "model_name_or_path: google/gemma-3-4b-pt\n",
        "trust_remote_code: true\n",
        "\n",
        "### method\n",
        "stage: pt\n",
        "do_train: true\n",
        "finetuning_type: full\n",
        "\n",
        "### deepspeed\n",
        "deepspeed: /mnt/cluster_storage/ds_z3_config.json # path to the DeepSpeed config\n",
        "\n",
        "### dataset\n",
        "dataset: my_cpt_c4\n",
        "dataset_dir: /mnt/cluster_storage\n",
        "\n",
        "template: gemma\n",
        "cutoff_len: 512\n",
        "max_samples: 1000\n",
        "overwrite_cache: true\n",
        "preprocessing_num_workers: 16\n",
        "\n",
        "### output\n",
        "output_dir: gemma3_4b_full_cpt\n",
        "logging_steps: 2\n",
        "save_steps: 50\n",
        "plot_loss: true\n",
        "report_to: mlflow   # or none\n",
        "\n",
        "### train\n",
        "per_device_train_batch_size: 1 # Adjust this depending on your GPU memory and sequence length\n",
        "gradient_accumulation_steps: 2\n",
        "num_train_epochs: 2.0\n",
        "learning_rate: 1.0e-4\n",
        "bf16: true\n",
        "lr_scheduler_type: cosine\n",
        "warmup_ratio: 0.1\n",
        "ddp_timeout: 180000000\n",
        "\n",
        "### ray\n",
        "ray_run_name: gemma3_4b_full_cpt\n",
        "ray_storage_path: /mnt/cluster_storage/\n",
        "ray_num_workers: 4  # Number of GPUs to use\n",
        "resources_per_worker:\n",
        "  GPU: 1\n",
        "  # accelerator_type:L40S: 0.001            # Use this to simply specify a GPU type (may place GPUs on separate nodes).\n",
        "  anyscale/accelerator_shape:4xL40S: 0.001  # Prefer this for DeepSpeed so all 4 GPUs are on the same node.\n",
        "  # See https://docs.ray.io/en/master/ray-core/accelerator-types.html#accelerator-types for a full list of accelerator types.\n",
        "ray_init_kwargs:\n",
        "  runtime_env:\n",
        "    env_vars:\n",
        "      # If using gated models like google/gemma-3-4b-pt\n",
        "      HF_TOKEN: <your_huggingface_token>\n",
        "      # If hf_transfer is installed\n",
        "      HF_HUB_ENABLE_HF_TRANSFER: '1'\n",
        "      # If using mlflow for experiments tracking\n",
        "      MLFLOW_TRACKING_URI: \"https://<your_cloud_id>.cloud.databricks.com\"\n",
        "      MLFLOW_TRACKING_TOKEN: \"<mlflow_tracking_token>\"\n",
        "      MLFLOW_EXPERIMENT_NAME: \"/Users/<your_user_id>/experiment_name\"\n",
        "```\n",
        "\n",
        "**Note:**\n",
        "This configuration assumes `4xL40S` GPUs are available in your cloud environment. If not, you can substitute with `4xA100-40G` (or another supported accelerator with similar VRAM).\n",
        "\n",
        "Together, `stage: pt` and `finetuning_type: full` configure this run as full continued pre-training on this C4-based corpus, producing full model checkpoints rather than lightweight adapters.\n",
        "\n",
        "### DeepSpeed configuration\n",
        "DeepSpeed is an open-source deep-learning optimization library developed by Microsoft, aimed at enabling large-model training. Higher ZeRO stages (1→3) and enabling CPU offload reduce GPU VRAM usage, but might cause slower training.\n",
        "\n",
        "To enable DeepSpeed, create a separate Deepspeed config in the **[storage shared across nodes](https://docs.anyscale.com/configuration/storage#shared)**. and reference it from your main training yaml config with:\n",
        "\n",
        "```yaml\n",
        "deepspeed: /mnt/cluster_storage/ds_z3_config.json\n",
        "```\n",
        "\n",
        "Below is a sample ZeRO-3 config:\n",
        "\n",
        "`ds_z3_config.json`\n",
        "```json\n",
        "{\n",
        "\"train_batch_size\": \"auto\",\n",
        "\"train_micro_batch_size_per_gpu\": \"auto\",\n",
        "\"gradient_accumulation_steps\": \"auto\",\n",
        "\"gradient_clipping\": \"auto\",\n",
        "\"zero_allow_untested_optimizer\": true,\n",
        "\"fp16\": {\n",
        "    \"enabled\": \"auto\",\n",
        "    \"loss_scale\": 0,\n",
        "    \"loss_scale_window\": 1000,\n",
        "    \"initial_scale_power\": 16,\n",
        "    \"hysteresis\": 2,\n",
        "    \"min_loss_scale\": 1\n",
        "},\n",
        "\"bf16\": {\n",
        "    \"enabled\": \"auto\"\n",
        "},\n",
        "\"zero_optimization\": {\n",
        "    \"stage\": 3,\n",
        "    \"overlap_comm\": false,\n",
        "    \"contiguous_gradients\": true,\n",
        "    \"sub_group_size\": 1e9,\n",
        "    \"reduce_bucket_size\": \"auto\",\n",
        "    \"stage3_prefetch_bucket_size\": \"auto\",\n",
        "    \"stage3_param_persistence_threshold\": \"auto\",\n",
        "    \"stage3_max_live_parameters\": 1e9,\n",
        "    \"stage3_max_reuse_distance\": 1e9,\n",
        "    \"stage3_gather_16bit_weights_on_model_save\": true\n",
        "}\n",
        "}\n",
        "```\n",
        "\n",
        "For a more detailed guide on acceleration and optimization methods including DeepSpeed on Ray, see [Speed and memory optimizations](https://docs.anyscale.com/llm/fine-tuning/speed-and-memory-optimizations).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%%bash\n",
        "# Create a copy of the DeepSpeed configuration file in /mnt/cluster_storage\n",
        "cp ../deepspeed-configs/ds_z3_config.json /mnt/cluster_storage/\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 4: Train and monitor\n",
        "\n",
        "**Note**: For gated models such as [`google/gemma-3-4b-pt`](https://huggingface.co/google/gemma-3-4b-pt), ensure that you accept the license agreement for the models on the Hugging Face site and set `HF_TOKEN` in the runtime environment. If you installed MLflow, configure its credentials. Otherwise, set `report_to: none` in `cpt_deepspeed.yaml` to avoid `api_token not set` errors.\n",
        "\n",
        "With all configurations in place, you can launch pre-training in one of two ways:\n",
        "\n",
        "### Option A: Run from a workspace (quickstart)\n",
        "\n",
        "The `USE_RAY=1` prefix tells LLaMA-Factory to run in distributed mode on the Ray cluster attached to your workspace.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%%bash\n",
        "USE_RAY=1 llamafactory-cli train ../train-configs/cpt_deepspeed.yaml\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Option B: Run as an Anyscale job (production)\n",
        "\n",
        "For longer or production runs, submit the training as an **Anyscale job**. Jobs run outside your interactive session for better stability, retries, and durable logs. You package LLaMA-Factory and other libraries in a container image and launch with a short job config. See [Run LLaMA-Factory as an Anyscale job](https://docs.anyscale.com/llm/fine-tuning/llamafactory-jobs) for the step-by-step guide.\n",
        "\n",
        "### Tracking with MLflow\n",
        "\n",
        "If you enabled MLflow logging (`report_to: mlflow` in your YAML), LLaMA-Factory logs metrics (loss, learning rate, etc.), parameters, and artifacts to your configured MLflow tracking server.\n",
        "\n",
        "**Example YAML snippet:**\n",
        "\n",
        "```yaml\n",
        "report_to: mlflow\n",
        "\n",
        "ray_init_kwargs:\n",
        "  runtime_env:\n",
        "    env_vars:\n",
        "      MLFLOW_TRACKING_URI: \"https://<your_cloud_id>.cloud.databricks.com\"\n",
        "      MLFLOW_TRACKING_TOKEN: \"<mlflow_tracking_token>\"\n",
        "      MLFLOW_EXPERIMENT_NAME: \"/Users/<your_user_id>/experiment_name\"\n",
        "```\n",
        "\n",
        "**MLFlow example**\n",
        "\n",
        "![MLflow](https://anyscale-public-materials.s3.us-west-2.amazonaws.com/llm-finetuning/llama-factory/3.2.4/mlflow.png)\n",
        "\n",
        "For a more detailed guide on tracking experiments with other tools such as Weights & Biases or MLflow, see [Observability and tracking](https://docs.anyscale.com/llm/fine-tuning/observability-and-tracking).\n",
        "\n",
        "## Step 5: Locate checkpoints\n",
        "\n",
        "Ray Train writes checkpoints under `ray_storage_path/ray_run_name`. In this example run, the path is: `/mnt/cluster_storage/gemma3_4b_full_cpt`. \n",
        "\n",
        "Inside, you see a **trainer session** directory named like:\n",
        "`TorchTrainer_8c6a5_00000_0_2025-09-09_09-53-45/`.\n",
        "\n",
        "- Ray Train creates `TorchTrainer_*` **when the trainer starts**; the suffix encodes a short run ID and the **start timestamp**.\n",
        "- Within that directory, Ray Train names checkpoints `checkpoint_000xxx/`, where the number is the saved ordered checkpoints.\n",
        "\n",
        "Control the save cadence with `save_strategy` and `save_steps`. For instructions on how to resume interrupted training with `resume_from_checkpoint` and more, see [Understand the artifacts directory](https://docs.anyscale.com/llm/fine-tuning/checkpointing#artifacts-directory).\n",
        "\n",
        "## Step 6: Export the model\n",
        "\n",
        "If you use LoRA, you can keep the base model and adapters separate for [multi-LoRA deployment](https://docs.anyscale.com/llm/serving/multi-lora) or [merge the adapters](https://docs.anyscale.com/llm/fine-tuning/checkpointing#merge-lora) into the base model for low-latency inference. \n",
        "\n",
        "For full fine-tuning or freeze-tuning, export the fine-tuned model directly.\n",
        "\n",
        "You may optionally apply [post-training quantization](https://docs.anyscale.com/llm/fine-tuning/checkpointing#ptq) on merged or full models before serving.\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.13.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}
